import redis
import time
import uuid
from flask import Flask, request, jsonify
from functools import wraps
from flask import jsonify
import gevent
from gevent import monkey
monkey.patch_all()
app = Flask(__name__)
# rds = redis.Redis.from_url("redis://:difyai123456@192.168.100.102:6380/0")
rds = redis.Redis.from_url("redis://:@127.0.0.1:6379/0")
class UserConcurrencyLimiter:
    LUA_SEMAPHORE = """
    -- KEYS[1] = semaphore key
    -- ARGV[1] = now
    -- ARGV[2] = timeout
    -- ARGV[3] = max concurrent
    -- ARGV[4] = token

    redis.call("ZREMRANGEBYSCORE", KEYS[1], 0, ARGV[1] - ARGV[2])
    local current = redis.call("ZCARD", KEYS[1])
    if tonumber(current) >= tonumber(ARGV[3]) then
        return 0
    end
    redis.call("ZADD", KEYS[1], ARGV[1], ARGV[4])
    return 1
    """
    def __init__(self, max_concurrent):
        self.max_concurrent = max_concurrent

    def get_allow(self, sema_key, token):
        allowed = rds.eval(
            self.LUA_SEMAPHORE,
            1, sema_key,
            str(time.time()),
            "10",
            str(self.max_concurrent),
            token
        )
        return allowed

    def user_concurrency_limit(self, func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            user_id = request.headers.get("X-USER-ID")
            if not user_id:
                return jsonify({"error": "缺少 X-USER_ID 请求头"}), 400
            sema_key = f"sema:user:{user_id}"
            token = str(uuid.uuid4())
            # limiter = UserConcurrencyLimiter(MAX_CONCURRENT)
            allowed = self.get_allow(sema_key, token)

            if allowed == 0:
                return jsonify({"error": "并发超限"}), 429
            try:
                resp = func(*args, **kwargs)
                return resp
            finally:
                rds.zrem(sema_key, token)
        return wrapper

LOCK_TIMEOUT = 10
MAX_CONCURRENT = 3
def user_concurrency_limit_2(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        user_id = request.headers.get("X-USER-ID")
        if not user_id:
            return jsonify({"error": "缺少 X-USER-ID 请求头"}), 400

        sema_key = f"sema:user:{user_id}"
        token = str(uuid.uuid4())
        now = time.time()

        pipe = rds.pipeline()
        try:
            pipe.zremrangebyscore(sema_key, 0, now - LOCK_TIMEOUT)
            pipe.zadd(sema_key, {token: now})
            pipe.zcard(sema_key)
            _, _, count = pipe.execute()

            if count > MAX_CONCURRENT:
                rds.zrem(sema_key, token)
                return jsonify({"error": "并发超限"}), 429

            resp = func(*args, **kwargs)
            return resp
        finally:
            rds.zrem(sema_key, token)
    return wrapper

limiter = UserConcurrencyLimiter(3)


@app.route("/api/t1")
@limiter.user_concurrency_limit
def t1():
    gevent.sleep(2)
    return jsonify({"status":200, "msg":"test user [user_concurrency_limit2]"})


@app.route("/api/t2")
@limiter.user_concurrency_limit
def api_t2():
    gevent.sleep(1)
    return  jsonify({"status":200, "msg":"test user"})

@app.route("/api/t3")
@limiter.user_concurrency_limit
def api_t3():
    gevent.sleep(1)
    return  jsonify({"status":200, "msg":"test user"})

@app.route("/api/fujg")
@limiter.user_concurrency_limit
def api_t5():
    gevent.sleep(1)
    return  jsonify({"status":200, "msg":"test user"})